import gradio as gr
from pages.removebg import portrait_matting,universal_matting
from modelscope.outputs import OutputKeys
from commonutils import  ImageUtils
import numpy as np


def inference(image, bk, model='universal', progress=gr.Progress()):
    # 一键抠图
    #progress(0, desc="正在抠图.....")
    #通用物体
    if model == 'universal':
      result = universal_matting(image)
    # 人像
    if model == 'portrait':
      result = portrait_matting(image)
    res_img=ImageUtils.imageOutput(result[OutputKeys.OUTPUT_IMG])
    #progress(0.5, desc="正在切换背景.....")
    alpha = res_img[:, :, 3:4] / 255.0
    w, h = image.size
    bk = bk.resize((w, h))
    # combine_img = Image.alpha_composite(Image.fromarray(res_img),bk)
    combine_img = image * alpha + bk * (1 - alpha)
    combine_img = combine_img.astype(np.uint8)
    return res_img, combine_img

def inference2(image, bk):
    return inference(image, bk,'portrait')

with gr.Blocks(css="#fixed_size_img {height: 240px;}") as changeBgApp:
    with gr.Tab("通用物体换背景"):
        with gr.Row():
            img_input1 = gr.Image(label="前图", type="pil", elem_id="fixed_size_img")
            img_output1 = gr.Image(label="抠图", type="pil", elem_id="fixed_size_img")
        with gr.Row():
            img_input2 = gr.Image(label="背景", type="pil", elem_id="fixed_size_img")
            img_output2 = gr.Image(label="新图", type="pil", elem_id="fixed_size_img")
        with gr.Row():
            btn_submit = gr.Button(value="一键抠图换背景")
        examples = [['./images/examples/chagebg_input2.jpg', './images/examples/chagebg_bk1.jpg'],
                  ['./images/examples/chagebg_input2.jpg', './images/examples/chagebg_bk2.jpg'],
                  ['./images/examples/chagebg_input3.jpg', './images/examples/chagebg_bk3.jpg']]
        examples = gr.Examples(examples=examples, inputs=[img_input1, img_input2])

    with gr.Tab("人像换背景"):
        with gr.Row():
            img_input3 = gr.Image(label="前图", type="pil", elem_id="fixed_size_img")
            img_output3 = gr.Image(label="抠图", type="pil", elem_id="fixed_size_img")
        with gr.Row():
            img_input4 = gr.Image(label="背景", type="pil", elem_id="fixed_size_img")
            img_output4 = gr.Image(label="新图", type="pil", elem_id="fixed_size_img")
        with gr.Row():
            btn_submit2 = gr.Button(value="一键抠图换背景")
        examples2 = [['./images/examples/chagebg_input2.jpg', './images/examples/chagebg_bk1.jpg'],
                  ['./images/examples/chagebg_input2.jpg', './images/examples/chagebg_bk2.jpg'],
                  ['./images/examples/chagebg_input3.jpg', './images/examples/chagebg_bk3.jpg']]
        examples2 = gr.Examples(examples=examples2, inputs=[img_input3, img_input4])

    btn_submit.click(inference, inputs=[img_input1, img_input2], outputs=[img_output1, img_output2])
    btn_submit2.click(inference2, inputs=[img_input3, img_input4], outputs=[img_output3, img_output4])



